package com.mmg.cosplay.service.impl;

import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import com.mmg.cosplay.entity.po.KnowledgeDocument;
import com.mmg.cosplay.mapper.KnowledgeDocumentMapper;
import com.mmg.cosplay.service.IKnowledgeService;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingSearchResult;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.qdrant.QdrantEmbeddingStore;
import io.qdrant.client.QdrantClient;
import io.qdrant.client.QdrantGrpcClient;
import io.qdrant.client.grpc.Collections.Distance;
import io.qdrant.client.grpc.Collections.VectorParams;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.apache.pdfbox.Loader;
import org.apache.pdfbox.pdmodel.PDDocument;
import org.apache.pdfbox.text.PDFTextStripper;
import org.apache.poi.xwpf.usermodel.XWPFDocument;
import org.apache.poi.xwpf.extractor.XWPFWordExtractor;
import org.apache.poi.hwpf.HWPFDocument;
import org.apache.poi.hwpf.extractor.WordExtractor;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;
import org.springframework.web.multipart.MultipartFile;

import java.io.File;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.time.LocalDateTime;
import java.util.List;
import java.util.UUID;
import java.util.stream.Collectors;

/**
 * 知识库服务实现类
 * 
 * @author fly
 * @since 2025-9-24
 */
@Slf4j
@Service
@RequiredArgsConstructor
public class KnowledgeServiceImpl extends ServiceImpl<KnowledgeDocumentMapper, KnowledgeDocument> implements IKnowledgeService {

    private final QdrantEmbeddingStore.Builder qdrantBuilder;
    private final EmbeddingModel embeddingModel;
    private QdrantClient qdrantClient;

    @Value("${file.upload.path}")
    private String uploadPath;
    
    @Value("${qdrant.host:localhost}")
    private String qdrantHost;
    
    @Value("${qdrant.port:6334}")
    private Integer qdrantPort;
    
    @Value("${qdrant.api-key:}")
    private String apiKey;

    /**
     * 获取角色专用的collection名称
     */
    private String getRoleCollectionName(Long roleId) {
        return "role_knowledge_" + roleId;
    }

    /**
     * 获取角色专用的向量存储
     */
    private EmbeddingStore<TextSegment> getRoleEmbeddingStore(Long roleId) {
        String collectionName = getRoleCollectionName(roleId);
        
        // 确保collection存在
        ensureCollectionExists(collectionName);
        
        return qdrantBuilder
                .collectionName(collectionName)
                .build();
    }
    
    /**
     * 确保collection存在，如果不存在则创建
     */
    private void ensureCollectionExists(String collectionName) {
        try {
            // 初始化Qdrant客户端（如果还没有初始化）
            if (qdrantClient == null) {
                QdrantGrpcClient.Builder clientBuilder = QdrantGrpcClient.newBuilder(qdrantHost, qdrantPort, false);
                if (apiKey != null && !apiKey.trim().isEmpty()) {
                    clientBuilder.withApiKey(apiKey);
                }
                qdrantClient = new QdrantClient(clientBuilder.build());
            }
            
            // 检查collection是否存在
            try {
                qdrantClient.getCollectionInfoAsync(collectionName).get();
                log.debug("Collection {} already exists", collectionName);
            } catch (Exception e) {
                // Collection不存在，创建它
                log.info("Creating collection: {}", collectionName);
                
                // 创建collection，向量维度设置为1024（text-embedding-v3的维度）
                VectorParams vectorParams = VectorParams.newBuilder()
                        .setSize(1024)  // 根据你使用的embedding模型调整维度
                        .setDistance(Distance.Cosine)
                        .build();
                
                qdrantClient.createCollectionAsync(collectionName, vectorParams).get();
                log.info("Successfully created collection: {}", collectionName);
            }
        } catch (Exception e) {
            log.error("Failed to ensure collection exists: {}", collectionName, e);
            throw new RuntimeException("Failed to create collection: " + e.getMessage());
        }
    }

    /**
     * 文档分割器
     */
    private DocumentSplitter getDocumentSplitter() {
        return DocumentSplitters.recursive(500, 50);
    }

    @Override
    @Transactional
    public Long addTextKnowledge(Long roleId, String title, String content) {
        try {
            // 1. 保存文档记录到数据库
            KnowledgeDocument document = new KnowledgeDocument()
                    .setRoleId(roleId)
                    .setTitle(title)
                    .setContent(content)
                    .setDocumentType("text")
                    .setFileSize((long) content.getBytes(StandardCharsets.UTF_8).length)
                    .setStatus(0) // 处理中
                    .setCreateTime(LocalDateTime.now());
            
            this.save(document);

            // 2. 处理文档并存储到向量数据库
            processAndStoreDocument(document, content);

            return document.getId();
        } catch (Exception e) {
            log.error("添加文本知识失败: roleId={}, title={}", roleId, title, e);
            throw new RuntimeException("添加文本知识失败: " + e.getMessage());
        }
    }

    @Override
    @Transactional
    public Long addFileKnowledge(Long roleId, MultipartFile file) {
        try {
            // 1. 获取文件信息
            String fileName = file.getOriginalFilename();
            String fileExtension = getFileExtension(fileName).toLowerCase();
            
            // 2. 根据文件类型解析内容
            String content = extractTextFromFile(file, fileExtension);
    
            // 3. 确保上传目录存在
            File uploadDir = new File(uploadPath);
            if (!uploadDir.exists()) {
                uploadDir.mkdirs();
            }
    
            // 4. 保存文件
            String filePath = uploadPath + "/" + UUID.randomUUID() + "_" + fileName;
            File targetFile = new File(filePath);
            file.transferTo(targetFile);
    
            // 5. 保存文档记录到数据库
            KnowledgeDocument document = new KnowledgeDocument()
                    .setRoleId(roleId)
                    .setTitle(fileName)
                    .setContent(content)
                    .setDocumentType(fileExtension)
                    .setFilePath(filePath)
                    .setFileSize(file.getSize())
                    .setStatus(0) // 处理中
                    .setCreateTime(LocalDateTime.now());
    
            this.save(document);
    
            // 6. 处理文档并存储到向量数据库
            processAndStoreDocument(document, content);
    
            return document.getId();
        } catch (IOException e) {
            log.error("添加文件知识失败: roleId={}, fileName={}", roleId, file.getOriginalFilename(), e);
            throw new RuntimeException("添加文件知识失败: " + e.getMessage());
        }
    }

    /**
     * 根据文件类型提取文本内容
     */
    private String extractTextFromFile(MultipartFile file, String fileExtension) throws IOException {
        switch (fileExtension) {
            case "pdf":
                return extractTextFromPDF(file);
            case "docx":
                return extractTextFromDocx(file);
            case "doc":
                return extractTextFromDoc(file);
            case "txt":
                return new String(file.getBytes(), StandardCharsets.UTF_8);
            default:
                // 对于其他文件类型，尝试作为文本文件读取
                return new String(file.getBytes(), StandardCharsets.UTF_8);
        }
    }

    /**
     * 从PDF文件提取文本
     */
    private String extractTextFromPDF(MultipartFile file) throws IOException {
        try (PDDocument document = Loader.loadPDF(file.getInputStream().readAllBytes())) {
            PDFTextStripper stripper = new PDFTextStripper();
            return stripper.getText(document);
        }
    }

    /**
     * 从DOCX文件提取文本
     */
    private String extractTextFromDocx(MultipartFile file) throws IOException {
        try (XWPFDocument document = new XWPFDocument(file.getInputStream());
             XWPFWordExtractor extractor = new XWPFWordExtractor(document)) {
            return extractor.getText();
        }
    }

    /**
     * 从DOC文件提取文本
     */
    private String extractTextFromDoc(MultipartFile file) throws IOException {
        try (HWPFDocument document = new HWPFDocument(file.getInputStream());
             WordExtractor extractor = new WordExtractor(document)) {
            return extractor.getText();
        }
    }

    /**
     * 处理文档并存储到向量数据库
     */
    private void processAndStoreDocument(KnowledgeDocument document, String content) {
        try {
            // 1. 创建文档对象
            Document doc = Document.from(content);

            // 2. 分割文档
            DocumentSplitter splitter = getDocumentSplitter();
            List<TextSegment> segments = splitter.split(doc);

            // 3. 为每个片段添加元数据
            for (int i = 0; i < segments.size(); i++) {
                TextSegment segment = segments.get(i);
                segment.metadata().put("documentId", document.getId().toString());
                segment.metadata().put("roleId", document.getRoleId().toString());
                segment.metadata().put("title", document.getTitle());
                segment.metadata().put("segmentIndex", String.valueOf(i));
            }

            // 4. 生成向量并存储
            EmbeddingStore<TextSegment> embeddingStore = getRoleEmbeddingStore(document.getRoleId());
            List<Embedding> embeddings = embeddingModel.embedAll(segments).content();
            List<String> vectorIds = embeddingStore.addAll(embeddings, segments);

            // 5. 更新文档状态
            document.setVectorId(String.join(",", vectorIds))
                    .setStatus(1) // 处理成功
                    .setUpdateTime(LocalDateTime.now());
            this.updateById(document);

            log.info("文档处理成功: documentId={}, segmentCount={}", document.getId(), segments.size());
        } catch (Exception e) {
            // 更新文档状态为失败
            document.setStatus(-1)
                    .setErrorMessage(e.getMessage())
                    .setUpdateTime(LocalDateTime.now());
            this.updateById(document);
            
            log.error("文档处理失败: documentId={}", document.getId(), e);
            throw new RuntimeException("文档处理失败: " + e.getMessage());
        }
    }

    @Override
    @Transactional
    public void deleteKnowledge(Long documentId, Long roleId) {
        try {
            // 1. 查询文档
            KnowledgeDocument document = this.getById(documentId);
            if (document == null || !document.getRoleId().equals(roleId)) {
                throw new RuntimeException("文档不存在或无权限删除");
            }

            // 2. 从向量数据库删除
            if (document.getVectorId() != null) {
                EmbeddingStore<TextSegment> embeddingStore = getRoleEmbeddingStore(roleId);
                String[] vectorIds = document.getVectorId().split(",");
                for (String vectorId : vectorIds) {
                    embeddingStore.remove(vectorId);
                }
            }

            // 3. 删除文件（如果存在）
            if (document.getFilePath() != null) {
                java.io.File file = new java.io.File(document.getFilePath());
                if (file.exists()) {
                    file.delete();
                }
            }

            // 4. 逻辑删除数据库记录
            document.setIsDeleted(1).setUpdateTime(LocalDateTime.now());
            this.updateById(document);

            log.info("删除知识文档成功: documentId={}, roleId={}", documentId, roleId);
        } catch (Exception e) {
            log.error("删除知识文档失败: documentId={}, roleId={}", documentId, roleId, e);
            throw new RuntimeException("删除知识文档失败: " + e.getMessage());
        }
    }

    @Override
    public List<KnowledgeDocument> getKnowledgeByRole(Long roleId) {
        LambdaQueryWrapper<KnowledgeDocument> wrapper = new LambdaQueryWrapper<>();
        wrapper.eq(KnowledgeDocument::getRoleId, roleId)
                .eq(KnowledgeDocument::getIsDeleted, 0)
                .orderByDesc(KnowledgeDocument::getCreateTime);
        return this.list(wrapper);
    }

    @Override
    public List<TextSegment> searchKnowledge(Long roleId, String query, int maxResults) {
        try {
            // 1. 生成查询向量
            Embedding queryEmbedding = embeddingModel.embed(query).content();

            // 2. 在角色专用的向量存储中搜索
            EmbeddingStore<TextSegment> embeddingStore = getRoleEmbeddingStore(roleId);

            // 3. 构建搜索请求
            EmbeddingSearchRequest searchRequest = EmbeddingSearchRequest.builder()
                    .queryEmbedding(queryEmbedding)
                    .maxResults(maxResults)
                    .minScore(0.7) // 设置最小相似度阈值
                    .build();

            // 4. 执行搜索
            EmbeddingSearchResult<TextSegment> searchResult = embeddingStore.search(searchRequest);

            // 5. 提取文本片段
            return searchResult.matches().stream()
                    .map(match -> match.embedded())
                    .collect(Collectors.toList());
        } catch (Exception e) {
            log.error("搜索知识库失败: roleId={}, query={}", roleId, query, e);
            throw new RuntimeException("搜索知识库失败: " + e.getMessage());
        }
    }

    @Override
    @Transactional
    public void clearRoleKnowledge(Long roleId) {
        try {
            // 1. 查询所有文档
            List<KnowledgeDocument> documents = getKnowledgeByRole(roleId);

            // 2. 删除向量数据库中的collection
            // 注意：这里需要根据Qdrant的API来删除整个collection
            // 由于langchain4j可能没有直接的删除collection方法，我们逐个删除向量
            EmbeddingStore<TextSegment> embeddingStore = getRoleEmbeddingStore(roleId);
            for (KnowledgeDocument document : documents) {
                if (document.getVectorId() != null) {
                    String[] vectorIds = document.getVectorId().split(",");
                    for (String vectorId : vectorIds) {
                        embeddingStore.remove(vectorId);
                    }
                }
            }

            // 3. 删除文件
            for (KnowledgeDocument document : documents) {
                if (document.getFilePath() != null) {
                    java.io.File file = new java.io.File(document.getFilePath());
                    if (file.exists()) {
                        file.delete();
                    }
                }
            }

            // 4. 逻辑删除数据库记录
            LambdaQueryWrapper<KnowledgeDocument> wrapper = new LambdaQueryWrapper<>();
            wrapper.eq(KnowledgeDocument::getRoleId, roleId)
                    .eq(KnowledgeDocument::getIsDeleted, 0);
            
            KnowledgeDocument updateEntity = new KnowledgeDocument()
                    .setIsDeleted(1)
                    .setUpdateTime(LocalDateTime.now());
            this.update(updateEntity, wrapper);

            log.info("清空角色知识库成功: roleId={}", roleId);
        } catch (Exception e) {
            log.error("清空角色知识库失败: roleId={}", roleId, e);
            throw new RuntimeException("清空角色知识库失败: " + e.getMessage());
        }
    }

    /**
     * 获取文件扩展名
     */
    private String getFileExtension(String fileName) {
        if (fileName == null || fileName.lastIndexOf('.') == -1) {
            return "unknown";
        }
        return fileName.substring(fileName.lastIndexOf('.') + 1).toLowerCase();
    }
}